from pydantic import BaseModel
from typing import Optional, List, Any


class AppKey(BaseModel):
    app_key: str


class TextToImageRequest(BaseModel):
    app_key: str
    task_id: str
    prompt: str
    prompt_args: dict
    negative_prompt: str
    width: int
    height: int
    batch_size: int
    steps: int
    cfg_scale: float


class TextToImageResponse(BaseModel):
    success: bool
    error_code: int = None
    error_msg: str = None
    list: list = None


class ImageToImageRequest(BaseModel):
    app_key: str
    task_id: str
    prompt: str
    prompt_args: dict
    negative_prompt: str
    width: int
    height: int
    batch_size: int
    steps: int
    cfg_scale: float
    from_image: str
    mask_image: str


class ImageToImageResponse(BaseModel):
    success: bool
    error_code: int = None
    error_msg: str = None
    list: list = None


class BackgroundGenerationRequest(BaseModel):
    app_key: str
    task_id: str
    prompt: str
    prompt_args: dict
    negative_prompt: str
    width: int
    height: int
    batch_size: int
    steps: int
    cfg_scale: float


class BackgroundGenerationResponse(BaseModel):
    success: bool
    error_code: int = None
    error_msg: str = None
    list: list = None


class RemoveBackgroundRequest(BaseModel):
    app_key: str
    task_id: str
    input_image: str


class RemoveBackgroundResponse(BaseModel):
    success: bool
    error_code: int = None
    error_msg: str = None
    data: dict = None


class TextGenerationRequest(BaseModel):
    app_key: str
    input_txt: str


class TextGenerationResponse(BaseModel):
    success: bool
    error_code: int = None
    error_msg: str = None
    data: dict = None


class ProgressQueryRequest(BaseModel):
    app_key: str
    task_id: str


class ProgressQueryResponse(BaseModel):
    success: bool
    error_code: int = None
    error_msg: str = None
    data: dict = None
